package gateway.config;

import lombok.AllArgsConstructor;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.session.CookieWebSessionIdResolver;
import org.springframework.web.server.session.HeaderWebSessionIdResolver;
import org.springframework.web.server.session.WebSessionIdResolver;

import java.util.List;

@Component
public class WebSessionIdResolverBeanPostProcessor implements BeanPostProcessor {
    public static final String HEADER_X_AUTH_TOKEN = "X-Auth-Token";

    public static final String HEADER_AUTHENTICATION_INFO = "Authentication-Info";

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (bean instanceof WebSessionIdResolver) {
            // https://docs.spring.io/spring-session/reference/guides/boot-webflux-custom-cookie.html
            HeaderWebSessionIdResolver resolver = new HeaderWebSessionIdResolver();
            resolver.setHeaderName(HEADER_X_AUTH_TOKEN);
            return new CombinationWebSessionIdResolver((CookieWebSessionIdResolver) bean, resolver);
        }
        return bean;
    }

    @AllArgsConstructor
    static class CombinationWebSessionIdResolver implements WebSessionIdResolver {
        CookieWebSessionIdResolver cookieWebSessionIdResolver;
        HeaderWebSessionIdResolver headerWebSessionIdResolver;

        @Override
        public List<String> resolveSessionIds(ServerWebExchange exchange) {
            List<String> sessionIds = headerWebSessionIdResolver.resolveSessionIds(exchange);
            if (sessionIds.isEmpty()) sessionIds = cookieWebSessionIdResolver.resolveSessionIds(exchange);
            return sessionIds;
        }

        @Override
        public void setSessionId(ServerWebExchange exchange, String sessionId) {
            headerWebSessionIdResolver.setSessionId(exchange, sessionId);
            cookieWebSessionIdResolver.setSessionId(exchange, sessionId);
        }

        @Override
        public void expireSession(ServerWebExchange exchange) {
            headerWebSessionIdResolver.expireSession(exchange);
            cookieWebSessionIdResolver.expireSession(exchange);
        }
    }
}